import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import solve_ivp, odeint
from scipy.fft import fft, fftfreq
import numpy.linalg as la
from matplotlib.animation import FuncAnimation
from mpl_toolkits.mplot3d import Axes3D
from sklearn.preprocessing import MinMaxScaler

# 设置随机种子和绘图风格
np.random.seed(42)
plt.style.use('dark_background')
# 中文字体
plt.rcParams['font.sans-serif'] = ['WenQuanYi Zen Hei']  # 使用黑体
# 解决负号显示问题
plt.rcParams['axes.unicode_minus'] = False  
class SSTMS_Simulator:
    def __init__(self):
        # 基本常数
        self.hbar = 1.0  # 简化普朗克常数
        self.c = 1.0     # 简化光速
        
        # 理论参数
        self.xi = 0.01   # SQ-WQ耦合常数 ξ
        self.M_cy = 10.0 # CY紧致化能标
        self.G_N = 1.0 / (self.M_cy**2)  # 有效牛顿常数
        
    def spontaneous_symmetry_breaking(self, phi_range=(-3, 3), num_points=1000):
        """
        模拟对称性自发破缺过程 - 墨西哥帽势
        """
        phi = np.linspace(phi_range[0], phi_range[1], num_points)
        
        # 墨西哥帽势: V(φ) = -μ²φ² + λφ⁴
        mu, lam = 1.0, 0.5
        V = -mu**2 * phi**2 + lam * phi**4
        
        # 寻找稳定真空
        vacuum_indices = np.where(V == np.min(V))[0]
        vacuum_phi = phi[vacuum_indices]
        
        fig, ax = plt.subplots(figsize=(10, 6))
        ax.plot(phi, V, 'cyan', linewidth=2, label='V(Φ) = -μ²Φ² + λΦ⁴')
        ax.scatter(vacuum_phi, np.min(V)*np.ones_like(vacuum_phi), 
                  color='red', s=100, zorder=5, label='破缺真空')
        ax.axhline(0, color='white', linestyle='--', alpha=0.3)
        ax.set_xlabel('原始场 Φ', fontsize=12)
        ax.set_ylabel('势能 V(Φ)', fontsize=12)
        ax.set_title('时空-物质对称性自发破缺', fontsize=14)
        ax.legend()
        ax.grid(True, alpha=0.3)
        
        return fig, vacuum_phi
    
    def quantum_tunneling(self, barrier_height=5.0, barrier_width=2.0):
        """
        模拟量子隧穿效应 - WQ引导的隧穿
        """
        x = np.linspace(-5, 5, 1000)
        
        # 方形势垒
        V = np.zeros_like(x)
        barrier_region = (x > -barrier_width/2) & (x < barrier_width/2)
        V[barrier_region] = barrier_height
        
        # 薛定谔方程解（简化）
        k = np.sqrt(2*(barrier_height - 1))  # 假设粒子能量为1
        transmission = np.exp(-2 * k * barrier_width)
        
        # 波函数形状（定性展示）
        psi = np.exp(-(x + 3)**2)  # 入射高斯波包
        psi_tunneled = transmission * np.exp(-(x - 3)**2)  # 隧穿部分
        
        fig, ax = plt.subplots(figsize=(12, 6))
        ax.plot(x, V, 'red', linewidth=2, label='势垒 V(x)')
        ax.plot(x, np.abs(psi)**2, 'cyan', linewidth=2, label='入射波 |ψ|²')
        ax.plot(x, np.abs(psi_tunneled)**2, 'yellow', linewidth=2, 
                label=f'隧穿波 |ψ|² (T={transmission:.3f})')
        
        ax.fill_between(x, V, alpha=0.3, color='red')
        ax.set_xlabel('位置 x', fontsize=12)
        ax.set_ylabel('概率密度 |ψ|² / 势能 V', fontsize=12)
        ax.set_title('WQ引导的量子隧穿效应', fontsize=14)
        ax.legend()
        ax.grid(True, alpha=0.3)
        
        return fig, transmission
    
    def entanglement_dynamics(self, separation=3.0, time_steps=100):
        """
        模拟纠缠动力学 - WQ场中的关联传播
        """
        t = np.linspace(0, 5, time_steps)
        
        # 两个纠缠粒子的位置
        x1, x2 = -separation/2, separation/2
        
        # WQ场关联函数随时间演化（简化模型）
        correlation = np.exp(-(t[:, None] - np.abs(np.linspace(-5, 5, 100) - x1)/self.c)**2)
        correlation += np.exp(-(t[:, None] - np.abs(np.linspace(-5, 5, 100) - x2)/self.c)**2)
        
        fig, ax = plt.subplots(figsize=(12, 8))
        im = ax.imshow(correlation.T, extent=[0, 5, -5, 5], aspect='auto', 
                      cmap='viridis', origin='lower')
        
        ax.axvline(x1, color='red', linestyle='--', label='粒子A位置')
        ax.axvline(x2, color='blue', linestyle='--', label='粒子B位置')
        ax.set_xlabel('时间 t', fontsize=12)
        ax.set_ylabel('空间位置 x', fontsize=12)
        ax.set_title('WQ场中的纠缠关联传播', fontsize=14)
        ax.legend()
        fig.colorbar(im, label='关联强度')
        
        return fig, correlation
    
    def casimir_effect(self, plate_separation=1.0, max_mode=50):
        """
        模拟卡西米尔效应 - WQ模式限制
        """
        d = plate_separation
        z = np.linspace(0.1, 3, 100)
        
        # 卡西米尔力计算公式
        F = -np.pi**2 * self.hbar * self.c / (240 * z**4)
        
        # 模式数量对比
        modes_free = max_mode  # 自由空间模式数
        modes_confined = int(max_mode * d / np.max(z))  # 受限空间模式数
        
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))
        
        # 力随距离变化
        ax1.plot(z, F, 'yellow', linewidth=2)
        ax1.set_xlabel('板间距 d', fontsize=12)
        ax1.set_ylabel('卡西米尔力 F', fontsize=12)
        ax1.set_title('卡西米尔力 vs 板间距', fontsize=14)
        ax1.grid(True, alpha=0.3)
        
        # 模式数量对比
        labels = ['自由空间', '受限空间']
        mode_counts = [modes_free, modes_confined]
        colors = ['cyan', 'red']
        ax2.bar(labels, mode_counts, color=colors, alpha=0.7)
        ax2.set_ylabel('WQ模式数量', fontsize=12)
        ax2.set_title('WQ模式限制效应', fontsize=14)
        
        return fig, F
    
    def gravitational_modification(self, r_range=(0.1, 2.0)):
        """
        模拟微小尺度的引力修正
        """
        r = np.linspace(r_range[0], r_range[1], 100)
        
        # 牛顿引力
        F_newton = -self.G_N / r**2
        
        # SST-MS理论修正（振荡修正项）
        F_sstms = F_newton * (1 + self.xi * np.sin(2 * np.pi * self.M_cy * r))
        
        fig, ax = plt.subplots(figsize=(12, 6))
        ax.plot(r, -F_newton, 'red', linewidth=2, label='牛顿引力')
        ax.plot(r, -F_sstms, 'cyan', linewidth=2, label='SST-MS理论修正')
        ax.fill_between(r, -F_newton, -F_sstms, alpha=0.3, color='gray')
        
        ax.set_xlabel('距离 r', fontsize=12)
        ax.set_ylabel('引力强度 |F|', fontsize=12)
        ax.set_title('微小尺度的引力修正（振荡效应）', fontsize=14)
        ax.legend()
        ax.grid(True, alpha=0.3)
        
        return fig, F_sstms

# 创建模拟器实例并运行模拟
simulator = SSTMS_Simulator()

print("开始SST-MS理论数值模拟...")
print("=" * 50)

# 1. 对称性自发破缺模拟
fig1, vacuum = simulator.spontaneous_symmetry_breaking()
plt.savefig('symmetry_breaking.png', dpi=300, bbox_inches='tight')
print(f"对称性破缺完成，真空期望值: {vacuum}")

# 2. 量子隧穿模拟
fig2, transmission = simulator.quantum_tunneling()
plt.savefig('quantum_tunneling.png', dpi=300, bbox_inches='tight')
print(f"量子隧穿模拟完成，透射系数: {transmission:.4f}")

# 3. 纠缠动力学模拟
fig3, correlation = simulator.entanglement_dynamics()
plt.savefig('entanglement.png', dpi=300, bbox_inches='tight')
print("纠缠动力学模拟完成")

# 4. 卡西米尔效应模拟
fig4, casimir_force = simulator.casimir_effect()
plt.savefig('casimir_effect.png', dpi=300, bbox_inches='tight')
print("卡西米尔效应模拟完成")

# 5. 引力修正模拟
fig5, modified_gravity = simulator.gravitational_modification()
plt.savefig('gravity_modification.png', dpi=300, bbox_inches='tight')
print("引力修正模拟完成")

print("=" * 50)
print("所有模拟完成！图像已保存。")

# 显示所有图像
plt.show()
